#!/usr/bin/env python
# -*- coding:utf-8 -*-
import os
import sys
import signal
import time,datetime
import subprocess
import traceback
import argparse
import yaml
import shutil
from lib.pssh.clients import ParallelSSHClient
from gevent import joinall

# 全局变量
client = None
dump_stat = True
flag = None

class Colors:
    FAIL = '\033[91m'
    GREEN = '\033[92m'
    WARNING = '\033[93m'
    CYAN = '\033[36m' 
    ENDC = '\033[0m'



def INFO(text):
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"{timestamp}-{Colors.CYAN}[INFO] {text}{Colors.ENDC}")

def SUCCESS(text):
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"{timestamp}-{Colors.GREEN}[SUCCESS] {text}{Colors.ENDC}")

def ERROR(text):
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"{timestamp}-{Colors.FAIL}[ERROR] {text}{Colors.ENDC}")

def WARN(text):
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"{timestamp}-{Colors.WARNING}[WARN] {text}{Colors.ENDC}")

def logo():
    print("                       _ _      _        _                       ")
    print(" _ __   __ _ _ __ __ _| | | ___| |    __| |_   _ _ __ ___  _ __  ")
    print("| '_ \ / _` | '__/ _` | | |/ _ \ |__ / _` | | | | '_ ` _ \| '_ \ ")
    print("| |_) | (_| | | | (_| | | |  __/ |__| (_| | |_| | | | | | | |_) |")
    print("| .__/ \__,_|_|  \__,_|_|_|\___|_|   \__,_|\__,_|_| |_| |_| .__/ ")
    print("|_|                                                       |_|    ")

# 处理Ctrl+C信号的逻辑
def signal_handler(signal, frame):
    global dump_stat,flag
    if flag == "dumping":
        sys.exit(1)
    if flag == "pcaping":
        dump_stat = False

# 参数解析器
def parser_args():
    # 创建解析器
    parser = argparse.ArgumentParser(description='open5gc parallel dump tools.',usage="paralleldump [-h] [-n NUM] [-f]")
    # 添加参数
    parser.add_argument("-n","--number",type=str,dest='num',default='test',help='用例编号(默认:test)')
    parser.add_argument("-f","--force",action='store_true',help='如果测试用例pcap包存在,强制删除')
    # 解析参数
    return parser.parse_args()

# 读取配置文件
def read_yaml_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as configfile:
        try:
            data = yaml.safe_load(configfile)
            return data
        except yaml.YAMLError as e:
            return None
# 合并pcap        
def merge_pcaps(copy_args,case_file):
    files = ""
    # 获取需要合并的pcap
    for index, item in enumerate(copy_args):
        for key,value in item.items():
            if key == 'local_file':
                if os.path.exists(value):
                    files = files + " " + value
                else:
                    ERROR("合并pcap失败:"+ value + "不存在")
                    return
    try:
        res = shutil.which('mergecap')
        if res == None:
            WARN("mergecap不存在,请使用yum install wireshark进行安装")
        else:
            command = f"mergecap -w {case_file} {files}"
            subprocess.run(command, shell=True)
    except Exception as e:
        ERROR("合并pcap失败")

if __name__ == '__main__':
    
    copy_args = []
    input_files = []
    hosts = []
    
    logo()

    # 解析参数
    args = parser_args()

    # 获取当前路径
    bundle_dir = os.path.abspath(os.path.dirname(__file__))
    exe_dir = bundle_dir.replace('lib','')
    # 配置文件
    config_file = os.path.join(exe_dir, 'config.yaml')
    # 保存pcap路径
    case_file_path = os.path.join(exe_dir, f'pcap/{args.num}')
    
    if not os.path.exists(config_file):
        ERROR("配置文件不存在")
        sys.exit(0)
    
    configdata = read_yaml_file(config_file)
    if configdata == None:
        ERROR("读取配置文件内容失败,请检查配置文件格式")
        sys.exit(0)
    else:
        user = configdata['config']['user']
        passwd = configdata['config']['passwd']
        port= configdata['config']['port']
        merge = configdata['operation']['merge']

   
    for index, item in enumerate(configdata['config']['host']):
        for key,value in item.items():
            hosts.append(value)
            copy_args.append({'remote_file': '/root/temp.pcap','local_file': f'{case_file_path}/{key}.pcap',})
            input_files.append(f'{case_file_path}/{key}.pcap')
    output_file = f'{case_file_path}/open5gc-{args.num}.pcap'

    INFO(f"当前测试用例编号为:{args.num}")

    if os.path.exists(case_file_path):
        if args.force:
            os.system(f'rm -rf {case_file_path}')
            os.makedirs(case_file_path)
        else:
            WARN("用例编号已经存在,使用-f/--force可强制删除.")
            sys.exit()
    else:
        os.makedirs(case_file_path)

    # 创建客户端：主机名，账号，密码，端口号
    client = ParallelSSHClient(hosts,user,passwd,port,num_retries=2,timeout=2)
    # 将信号处理器与SIGINT信号关联
    signal.signal(signal.SIGINT, signal_handler)
    try:
        INFO("正在尝试连接主机...")
        output = client.run_command('uname')
        for host_out in output:
            for line in host_out.stdout:
                if line == "Linux":
                    SUCCESS(f"主机{host_out.host}连接成功")
        INFO("正在进行抓包(使用CTRL+C停止抓包)...")
        flag = "dumping"
        client.run_command('tcpdump -i any -w /root/temp.pcap')
    except Exception as e:
        ERROR(f"主机连接失败:{e}")
        sys.exit(1)

    try:
        flag = "pcaping"
        while True:
            if dump_stat:
                time.sleep(0.5)
            else:
                client.run_command('kill -INT $(pidof tcpdump)') # 发送INT信号
                time.sleep(1)
                INFO("拷贝文件中...")
                cmds = client.copy_remote_file('%(remote_file)s', '%(local_file)s',copy_args=copy_args)
                joinall(cmds,raise_error=False)
                client.run_command('rm -rf /root/temp.pcap && exit 0')
                if merge: 
                    INFO(f"合并pcap中...")
                    merge_pcaps(copy_args,output_file)
                SUCCESS(f"文件已保存到:")
                os.system(f"tree -L 1 {case_file_path}")
                break
    except Exception as e:
        ERROR(e)